import os
import re

import cv2
import numpy as np
from get_coor import get_box
import plt_fit as pf
from ui_utils import r2_score


def getlocs(src, box, k, b, focus_losc):
    x,y,w,h=box
    I = src[y:y + h, x:x + w, :]
    src = 0.2989 * I[:, :, 0] + 0.587 * I[:, :, 1] + 0.114 * I[:, :, 2]
    II = np.uint8(src)
    num_bins = 256
    counts, _ = np.histogram(II, bins=num_bins, range=(0, 256))
    p = counts / np.sum(counts)
    omega = np.cumsum(p)
    mu = np.cumsum(p * np.arange(num_bins))
    mu_t = mu[-1]
    sigma_b_squared = (mu_t * omega - mu)**2 / (omega * (1 - omega))
    maxval = np.max(sigma_b_squared)
    idx = np.mean(np.where(sigma_b_squared == maxval))
    level = (idx - 1) / (num_bins - 1)
    _, bw = cv2.threshold(src, level * 255, 255, cv2.THRESH_BINARY)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    openbw = cv2.morphologyEx(bw.astype(np.uint8), cv2.MORPH_OPEN, kernel)
    stats = cv2.connectedComponentsWithStats(openbw, connectivity=8)[2]
    xx = stats[:, cv2.CC_STAT_LEFT][1]  # Assuming the second component is the one you want
    locs = k * xx + b
    L = focus_losc - locs
    return locs, L


def merge_intersecting_boxes(boxes):
    merged_boxes = []

    # 计算包含所有框的大框
    x_min = min(box[0] for box in boxes)
    y_min = min(box[1] for box in boxes)
    x_max = max(box[0] + box[2] for box in boxes)
    y_max = max(box[1] + box[3] for box in boxes)
    big_box = (x_min, y_min, x_max - x_min, y_max - y_min)

    # 返回大框和空的合并框列表
    return big_box, merged_boxes

def get_files(dir):
    img_path_list = [f for f in os.listdir(dir) if
                     f.startswith('Point') and f.endswith('.jpg')]  # 获取该文件夹中所有jpg格式的图像
    val_list=[]
    for p in img_path_list:
        # 使用正则表达式匹配_后.前的0或0.5
        match = re.search(r'_(\d+(\.\d+)?)\.', p)
        if match:
            val=match.group(1)
            val_list.append(float(val))
        else:
            raise ValueError('{0}文件名错误,无法提取位置i学那些'.format(p))
    return img_path_list,val_list


if __name__=='__main__':
    file_path = './20240531_113524'  # 图像文件夹路径
    img_path_list,locs=get_files(file_path)

    coors = []
    boxs = []
    for i, image_name in enumerate(img_path_list):  # 逐一读取图像
        item = cv2.imread(os.path.join(file_path, image_name))
        cneter, box, _ = get_box(item)
        coors.append(list(cneter))
        boxs.append(box)
    merge_box,_ = merge_intersecting_boxes(boxs)

    # 使用线性回归拟合数据
    matx = np.array(coors)
    arr_x = matx[:, 0]
    k, b = np.polyfit(matx[:, 0], locs, 1)
    y_true = locs
    y_pred = k * arr_x + b
    r2 = r2_score(y_true, y_pred)
    # 输出 R^2 值
    print("polyfit R^2 值:", r2)
    path = '20240531_113524\Point20_9.5.jpg'
    target = cv2.imread(path,cv2.IMREAD_COLOR)
    focus_losc = 0
    my_center, _, _ = get_box(target)
    x = my_center[0]
    print(k * x + b)

    draw_img=target.copy()
    x, y, w, h = merge_box
    cv2.rectangle(draw_img, (x, y), (x + w, y + h), (0, 255, 0), 2)

    pf.plot_image_and_r2(draw_img,arr_x,locs,k,b,r2)


